import asyncio
import json

from websockets.server import serve

from model.CNNModel import CNNModel
from model.LSTMModel import LSTMModel

cnn = CNNModel()
lstm = LSTMModel()
if __name__ == '__main__':

    async def handler(websocket):
        async for message in websocket:
            # 解析message格式
            try:
                decoder = json.JSONDecoder()
                obj = decoder.decode(message)
                if obj.get("command") == 'cnn-train':
                    cnn.setWS(websocket)
                    cnn.setCode(obj.get("code"))
                    cnn.setTime(obj.get("startTime"), obj.get("endTime"))
                    await  cnn.train()

                if obj.get("command") == 'cnn-draw':
                    cnn.setWS(websocket)
                    cnn.setCode(obj.get("code"))
                    cnn.setTime(obj.get("startTime"), obj.get("endTime"))
                    await  cnn.draw()

                if obj.get("command") == 'cnn-view-data':
                    cnn.setWS(websocket)
                    cnn.setCode(obj.get("code"))
                    cnn.setTime(obj.get("startTime"), obj.get("endTime"))
                    await  cnn.view()

                if obj.get("command") == 'lstm-train':
                    lstm.setWS(websocket)
                    lstm.setCode(obj.get("code"))
                    lstm.setTime(obj.get("startTime"), obj.get("endTime"))
                    await  lstm.train()

                if obj.get("command") == 'lstm-draw':
                    lstm.setWS(websocket)
                    lstm.setCode(obj.get("code"))
                    lstm.setTime(obj.get("startTime"), obj.get("endTime"))
                    await  lstm.draw()
                if obj.get("command") == 'lstm-view-data':
                    lstm.setWS(websocket)
                    lstm.setCode(obj.get("code"))
                    lstm.setTime(obj.get("startTime"), obj.get("endTime"))
                    await  lstm.view()
            except json.JSONDecodeError:
                obj = {
                    "code": 400,
                    "message": "json 格式化错误"
                }
                await  websocket.send(json.JSONEncoder().encode(obj))
            # except Exception:
            #     obj = {
            #         "code": 400,
            #         "message": "其他错误"
            #     }
            #     await  websocket.send(json.JSONEncoder().encode(obj))


    async def main():
        async with serve(handler, "localhost", 8888):
            await asyncio.Future()  # run forever


    asyncio.run(main())
